​NeurIPS 2022 | 知识蒸馏造成了样本不均衡问题？

©Paperweekly 原创 · 作者 | An.

单位 | 中科院自动化所

研究方向 | 计算机视觉、模型压缩

论文标题：

Respecting Transfer Gap in Knowledge Distillation

论文链接：

https://arxiv.org/pdf/2210.12787.pdf

引言

知识蒸馏（KD）[1] 是指将知识从教师模型转移到学生模型，以提高学生模型性能的一种模型压缩方法。一个普遍的直觉是，一个更强大的教师应该会教出一个表现更好的学生。然而，最近的研究发现，大模型不一定是好教师 [2]，早停或更小模型尺寸的教师反而能使学生获得更好的表现。

本文从域迁移的视角探究和分析了这一现象。作者认为 ground-truth（分类损失）和教师预测（蒸馏损失）的知识来自不同领域，其中存在着被过往方法忽视的转移差距。Ground-truth 所蕴含的知识所在的域可定义为“人类域（human domain）”，教师预测所蕴含的知识所在的域可定义为“机器域（machine domain）”。

通常来讲，数据集的构建会考虑类别样本的均衡问题，以 CIFAR100 为例，它包含 60000 张图像和 100 个类别，每个类别分别有 500 张训练图像和 100 张测试图像。

因此，人类域的知识在类别间是均衡的，即满足 （独立同分布）假设。而实验发现，教师预测的概率分布在类别间是不平衡的，无论是“强”还是“弱”的教师，由于转移差距的存在，教师预测尾部的类别的知识传递受阻。因此，机器域的知识在类别间是均衡的，即不满足 （独立同分布）假设。

原始 KD 忽视了这种转移差距，为分类损失和蒸馏损失分配了恒定权重。由于教师的知识是不平衡的，蒸馏损失的恒定样本权重将成为知识传递的瓶颈。

本文提出了一种简单有效的方法，即逆概率加权蒸馏（IPWD），来补充机器域中欠加权的训练样本。作者首先通过比较类感知预测和上下文感知预测来估计机器域的倾向得分 。然后，IPWD 将逆概率 作为蒸馏损失的样本权重，以加强不足样本的权重。通过这种方式，IPWD 生成了一个伪全体样本（pseudo-population）来处理不平衡的知识。

KD 的转移差距

2.1 教师知识的不平衡

以 CIFAR100 和 ImageNet 为例，作者统计了教师模型在训练样本中不同类别的预测概率和，例如，一张 dog 的样本的软标签为 ，则分别在 类和 类累加对应的数值。图 1 反映了 ground-truth 和教师预测间的分布差异，尽管教师模型是在平衡数据（蓝色虚线）上训练的，但其预测分布在不同温度上均存在不平衡问题。

▲ 图1. CIFAR100 和 ImageNet 上 ground-truth 和教师预测的分布差异

同时，作者还在 CIFAR100 上将类别按教师预测概率和排序分为 4 组，分别进行知识蒸馏。如表 1 所示，与普通训练相比，KD 在所有亚组中都取得了更好的表现。然而，前 25 个类带来的提升远高于后 25 个类别（平均 5.14% vs. 0.85%）。这证明了教师预测所蕴含的知识确实存在类别不平衡问题。

▲ 表1. 不同类别分组知识蒸馏的性能比较

2.2 因果推断视角的解释

作者从因果推断（causal inference）的角度解释了转移差距及其混淆效应。图 2 说明了图像 、训练数据 、教师参数 和教师输出 之间的因果关系。

总的来说， 和  的联合转移集是  和 的混淆因素。首先，训练集 和教师模型的转移集 共享相同图像集， 从图像集 中采样，即 是 的因。其次，教师 在 上训练， 由 和 计算，即 。因此， 和 是 的因。

值得注意的事，转移集是基于训练集 和教师模型 构建的，所以，我们将 和  的联合转移集 视为 和 的混淆因素。

▲ 图2. KD 的因果图

尽管在考虑上下文不变的类特定信息时 是平衡的，但它忽略了上下文信息（例如属性），这使得 在上下文上不平衡。如图 1 所示，这种不平衡的上下文导致不平衡的转移集 ，并进一步影响知识蒸馏的性能。

方法：IPWD

3.1 KD 的逆概率加权

从第 2 节的分析可知，人类域的 训练样本在机器域中不再是 。简单地将训练集假设为完美的转移集可能导致选择偏差：匹配“头部”知识的样本被过度表示，而匹配“尾部”知识的样本被低估。这将抑制“尾部”知识的转移。基于因果关系的理论 [3]，本文使用了逆概率加权（IPW）技术来克服转移差距导致的混淆效应。

简而言之，利用 IPW 生成一个伪全体样本，为欠表示的样本分配更大的权重，为过表示的样本分配较小的权重，以实现去偏蒸馏。IPW 将选择域倾向得分的倒数作为样本 的权重，新的目标函数为：

3.2 倾向得分

▲ 图3. 原始 KD 和 IPWD 的训练流程比较

本文提出了一种无监督的方式来估计机器域的倾向得分 ，利用人类域使用 CLS 训练的分类头作为参考，与 KD 训练的分类器的输出做比较，以确定样本是否在机器域的代表性不足。如图 3 所示，本文的方法训练了一个额外的分类头来计算 ，并根据标准差分别对 和 进行归一化，即，倾向得分的具体计算公式如下：

其中， 即 和 的交叉熵，。

IPWD 最终的目标函数可表示为：

实验

1. 图像分类-CIFAR100：如表 2 所示，IPWD 在绝大多数情况都取得了最好表现。由于本文的方法缓解了非 导致的转移差距，IPWD 在异构蒸馏中以更大的优势领先于其他 KD 方法。

▲ 表2.  CIFAR100 上不同蒸馏方法的对比

2. 图像分类-ImageNet：如表 3 所示，IPWD 在异构蒸馏中取得了远超其他 KD 方法的性能，不同网络结构的蒸馏在工程实践中也更加实用。

▲ 表3. ImageNet 上不同蒸馏方法的对比

3. 自蒸馏：从表 4 中可知，本文的方法可作为插件应用在 PS-KD [1] （目前最先进的自蒸馏方法）中，在不同架构下，IPWD 可带来 0.33~0.82%的性能提升。

▲ 表4. CIFAR100 上不同自蒸馏方法的对比

4. 消融实验：在“方法”部分，我们介绍了本文提出的两大改进（额外分类头和 logits 归一化）。表 5 证明了这两个改进的有效性，这两大改进都有助于降低方差，提高训练的稳定性。

▲ 表5. 消融实验（CLS head：额外分类头，logits norm.：logits 归一化）

参考文献

[1] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network.arXiv preprint arXiv:1503.02531, 2015.

[2] Jang Hyun Cho and Bharath Hariharan. On the efficacy of knowledge distillation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4794–4802, 2019.

[3] Paul R Rosenbaum and Donald B Rubin. The central role of the propensity score in observational studies for causal effects. Biometrika, 70(1):41–55, 1983.

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·

·